from langchain.chains import RetrievalQA #检索QA链，在文档上进行检索
from langchain_community.document_loaders import CSVLoader #文档加载器，采用csv格式存储

from langchain_community.vectorstores import DocArrayInMemorySearch #向量存储
import pandas as pd
#导入向量存储索引创建器
from langchain.indexes import VectorstoreIndexCreator
from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型
from openai import RateLimitError

from tool import get_azure_endpoint,get_api_key,get_api_version
import tenacity
from langchain_openai import AzureChatOpenAI
from langchain.output_parsers import StructuredOutputParser, ResponseSchema
from langchain.prompts import PromptTemplate

if __name__ == '__main__':
    file = './OutdoorClothingCatalog_1000.csv'
    # 使用langchain文档加载器对数据进行导入
    loader = CSVLoader(file_path=file)
    docs = loader.load()
    docs = docs[:100]
    # 查看单个文档，每个文档对应于CSV中的一行数据
    print(docs[0])

    # 初始化嵌入模型（这里使用OpenAI的嵌入模型，你也可以选择其他模型）
    embeddings = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
    )

    # 因为文档比较短了，所以这里不需要进行任何分块,可以直接进行向量表征 #使用初始化OpenAIEmbedding实例上的查询方法embed_query为文本创建向量表征
    embed = embeddings.embed_query("你好呀，我的名字叫小可爱")
    # 查看得到向量表征的长度
    print("\n\033[32m向量表征的长度: \033[0m \n", len(embed))
    # 每个元素都是不同的数字值，组合起来就是文本的向量表征
    print("\n\033[32m向量表征前5个元素: \033[0m \n", embed[:5])

    # 将刚才创建文本向量表征(embeddings)存储在向量存储(vector store)中 # 使用DocArrayInMemorySearch类的from_documents方法来实现
    # 该方法接受文档列表以及向量表征模型作为输入
    @tenacity.retry(
        wait=tenacity.wait_fixed(60),
        retry=tenacity.retry_if_exception_type(RateLimitError),
        stop=tenacity.stop_after_attempt(3)
    )
    def create_vector_store(docs, embeddings):
        return DocArrayInMemorySearch.from_documents(docs, embeddings)


    db = create_vector_store(docs, embeddings)
    query = "请推荐一件具有防晒功能的衬衫"
    # 使用上面的向量存储来查找与传入查询类似的文本，得到一个相似文档列表
    docs = db.similarity_search(query)
    print("\n\033[32m返回文档的个数: \033[0m \n", len(docs))
    print("\n\033[32m第一个文档: \033[0m \n", docs[0])

    # 导入大语言模型, 这里使用默认模型gpt-4o
    llm = AzureChatOpenAI(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        azure_deployment="gpt-4o",  # 重命名为 azure_deployment
        model_name="gpt-4o",
        openai_api_version=get_api_version(),  # 参数名不变
        openai_api_key=get_api_key(),
        openai_api_type="azure",
        temperature=0.0,
    )

    # 合并获得的相似文档内容
    qdocs = "".join([docs[i].page_content for i in range(len(docs))])
    # 将合并的相似文档内容后加上问题(question)输入到 `llm.call_as_llm`中 #这里问题是:以Markdown表格的方式列出所有具有防晒功能的衬衫并总结
    response = llm.call_as_llm(f"{qdocs}问题:请用markdown表格的方式列出所有具有防晒功能的衬衫，对每件衬衫描述进行总结")
    print(response)

    # 基于向量储存，创建检索器, stuff类型
    retriever = db.as_retriever()
    qa_stuff = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        verbose=True
    )
    # 创建一个查询并在此查询上运行链
    query = "请用markdown表格的方式列出所有具有防晒功能的衬衫，对每件衬衫描述进行总结"
    response = qa_stuff.run(query)
    print("\nstuff: \n", response)

    # map_reduce
    qa_map_reduce = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="map_reduce",
        retriever=retriever,
        verbose=True
    )
    # 创建一个查询并在此查询上运行链
    query = "请用markdown表格的方式列出所有具有防晒功能的衬衫，对每件衬衫描述进行总结"
    response = qa_map_reduce.run(query)
    print("\nmap_reduce: \n", response)



    # refine
    qa_refine = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="refine",
        retriever=retriever,
        verbose=True
    )
    # 创建一个查询并在此查询上运行链
    query = "请用markdown表格的方式列出所有具有防晒功能的衬衫，对每件衬衫描述进行总结"
    response = qa_refine.run(query)
    print("\nrefine: \n", response)